package com.future;

import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

/**
 * Description:
 * https://leetcode-cn.com/problems/encode-n-ary-tree-to-binary-tree/
 *
 * @author weiruibai.vendor
 * Date: 2022/2/18 10:45
 */
public class Solution_431 {

    public static void main(String[] args) {
        Node head = new Node(0);
        Node n1 = new Node(1);
        Node n2 = new Node(2);
        Node n3 = new Node(3);
        Node n4 = new Node(4);
        Node n5 = new Node(5);
        Node n6 = new Node(6);
        Node n7 = new Node(7);
        Node n8 = new Node(8);
        Node n9 = new Node(9);
        List<Node> children = new ArrayList<Node>() {{
            add(n1);
            add(n2);
            add(n3);
            add(n4);
        }};
        head.children = children;
        n1.children = new ArrayList<Node>() {{
            add(n5);
            add(n6);
        }};
        n2.children = new ArrayList<Node>() {{
            add(n7);
            add(n8);
            add(n9);
        }};
        Solution_431 encodeDecode = new Solution_431();
        TreeNode treeNode = encodeDecode.encode(head);
        Node node = encodeDecode.decode(treeNode);
        System.out.println(node);
    }

    /**
     * 参考 Solution_431.png
     *
     * @param root
     * @return
     */
    public TreeNode encode(Node root) {
        if (root == null) {
            return null;
        }
        TreeNode head = new TreeNode(root.val);
        head.left = en(root.children);
        return head;
    }

    private TreeNode en(List<Node> children) {
        if (children == null || children.isEmpty()) {
            return null;
        }
        TreeNode head = null;
        TreeNode curNode = null;
        for (Node node : children) {
            TreeNode cur = new TreeNode(node.val);
            if (head == null) {
                head = cur;
                curNode = head;
            } else {
                curNode.right = cur;
                curNode = cur;
            }
            curNode.left = en(node.children);
        }
        return head;
    }

    public Node decode(TreeNode root) {
        if (root == null) {
            return null;
        }
        return new Node(root.val, de(root.left));
    }

    private List<Node> de(TreeNode left) {
        List<Node> children = new ArrayList<>();
        while (left != null) {
            Node cur = new Node(left.val, de(left.left));
            children.add(cur);
            left = left.right;
        }
        return children;
    }


    static class Node {
        public int val;
        public List<Node> children;

        public Node() {
        }

        public Node(int _val) {
            val = _val;
        }

        public Node(int _val, List<Node> _children) {
            val = _val;
            children = _children;
        }


        @Override
        public String toString() {
            return "Node{" +
                    "val=" + val +
                    ", children=" + (children == null ? new ArrayList<Node>() : children).stream().map(obj -> obj.val).collect(Collectors.toList()) +
                    '}';
        }
    }

    private static class TreeNode {
        int val;
        TreeNode left;
        TreeNode right;

        TreeNode(int x) {
            val = x;
        }

        @Override
        public String toString() {
            return "TreeNode{" +
                    "val=" + val +
                    ", left=" + left.val +
                    ", right=" + right.val +
                    '}';
        }
    }
}
